import torch
import torchvision


class Basic(torch.nn.Module):
    def __init__(self, strType, intChannels):
        super().__init__()

        if strType == 'relu-conv-relu-conv':
            self.netMain = torch.nn.Sequential(
                torch.nn.PReLU(num_parameters=intChannels[0], init=0.25),
                torch.nn.Conv2d(in_channels=intChannels[0], out_channels=intChannels[1], kernel_size=3, stride=1, padding=1),
                torch.nn.PReLU(num_parameters=intChannels[1], init=0.25),
                torch.nn.Conv2d(in_channels=intChannels[1], out_channels=intChannels[2], kernel_size=3, stride=1, padding=1)
            )

        elif strType == 'conv-relu-conv':
            self.netMain = torch.nn.Sequential(
                torch.nn.Conv2d(in_channels=intChannels[0], out_channels=intChannels[1], kernel_size=3, stride=1, padding=1),
                torch.nn.PReLU(num_parameters=intChannels[1], init=0.25),
                torch.nn.Conv2d(in_channels=intChannels[1], out_channels=intChannels[2], kernel_size=3, stride=1, padding=1)
            )

        # end

        if intChannels[0] == intChannels[2]:
            self.netShortcut = None

        elif intChannels[0] != intChannels[2]:
            self.netShortcut = torch.nn.Conv2d(in_channels=intChannels[0], out_channels=intChannels[2], kernel_size=1, stride=1, padding=0)

        # end
    # end

    def forward(self, tenInput):
        if self.netShortcut is None:
            return self.netMain(tenInput) + tenInput

        elif self.netShortcut is not None:
            return self.netMain(tenInput) + self.netShortcut(tenInput)

        # end
    # end
# end

class Downsample(torch.nn.Module):
    def __init__(self, intChannels):
        super().__init__()

        self.netMain = torch.nn.Sequential(
            torch.nn.PReLU(num_parameters=intChannels[0], init=0.25),
            torch.nn.Conv2d(in_channels=intChannels[0], out_channels=intChannels[1], kernel_size=3, stride=2, padding=1),
            torch.nn.PReLU(num_parameters=intChannels[1], init=0.25),
            torch.nn.Conv2d(in_channels=intChannels[1], out_channels=intChannels[2], kernel_size=3, stride=1, padding=1)
        )
    # end

    def forward(self, tenInput):
        return self.netMain(tenInput)
    # end
# end

class Upsample(torch.nn.Module):
    def __init__(self, intChannels):
        super().__init__()

        self.netMain = torch.nn.Sequential(
            torch.nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
            torch.nn.PReLU(num_parameters=intChannels[0], init=0.25),
            torch.nn.Conv2d(in_channels=intChannels[0], out_channels=intChannels[1], kernel_size=3, stride=1, padding=1),
            torch.nn.PReLU(num_parameters=intChannels[1], init=0.25),
            torch.nn.Conv2d(in_channels=intChannels[1], out_channels=intChannels[2], kernel_size=3, stride=1, padding=1)
        )
    # end

    def forward(self, tenInput):
        return self.netMain(tenInput)
    # end
# end

class Refine(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.netImageOne = Basic('conv-relu-conv', [ 3, 24, 24 ])
        self.netImageTwo = Downsample([ 24, 48, 48 ])
        self.netImageThr = Downsample([ 48, 96, 96 ])

        self.netDisparityOne = Basic('conv-relu-conv', [ 1, 96, 96 ])
        self.netDisparityTwo = Upsample([ 192, 96, 96 ])
        self.netDisparityThr = Upsample([ 144, 48, 48 ])
        self.netDisparityFou = Basic('conv-relu-conv', [ 72, 24, 24 ])

        self.netRefine = Basic('conv-relu-conv', [ 24, 24, 1 ])
    # end

    def forward(self, tenImage, tenDisparity):
        tenMean = [ tenImage.mean([1, 2, 3], True), tenDisparity.mean([1, 2, 3], True) ]
        tenStd = [ tenImage.std([1, 2, 3], False, True), tenDisparity.std([1, 2, 3], False, True) ]

        tenImage = tenImage.clone()
        tenImage -= tenMean[0]
        tenImage /= tenStd[0] + 0.0000001

        tenDisparity = tenDisparity.clone()
        tenDisparity -= tenMean[1]
        tenDisparity /= tenStd[1] + 0.0000001

        tenImageOne = self.netImageOne(tenImage)
        tenImageTwo = self.netImageTwo(tenImageOne)
        tenImageThr = self.netImageThr(tenImageTwo)

        tenUpsample = self.netDisparityOne(tenDisparity)
        if tenUpsample.shape != tenImageThr.shape: tenUpsample = torch.nn.functional.interpolate(input=tenUpsample, size=(tenImageThr.shape[2], tenImageThr.shape[3]), mode='bilinear', align_corners=False) # not ideal
        tenUpsample = self.netDisparityTwo(torch.cat([ tenImageThr, tenUpsample ], 1)); tenImageThr = None
        if tenUpsample.shape != tenImageTwo.shape: tenUpsample = torch.nn.functional.interpolate(input=tenUpsample, size=(tenImageTwo.shape[2], tenImageTwo.shape[3]), mode='bilinear', align_corners=False) # not ideal
        tenUpsample = self.netDisparityThr(torch.cat([ tenImageTwo, tenUpsample ], 1)); tenImageTwo = None
        if tenUpsample.shape != tenImageOne.shape: tenUpsample = torch.nn.functional.interpolate(input=tenUpsample, size=(tenImageOne.shape[2], tenImageOne.shape[3]), mode='bilinear', align_corners=False) # not ideal
        tenUpsample = self.netDisparityFou(torch.cat([ tenImageOne, tenUpsample ], 1)); tenImageOne = None

        tenRefine = self.netRefine(tenUpsample)
        tenRefine *= tenStd[1] + 0.0000001
        tenRefine += tenMean[1]

        return torch.nn.functional.threshold(input=tenRefine, threshold=0.0, value=0.0)
    # end
# end



# netRefine = Refine().cuda().eval()
# netRefine.load_state_dict({ strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.hub.load_state_dict_from_url(url='http://content.sniklaus.com/kenburns/network-refinement.pytorch', file_name='kenburns-refinement').items() })

# def disparity_refinement(tenImage, tenDisparity):
#     return netRefine(tenImage, tenDisparity)
# end